# challenge 1

import pandas as pd
import numpy as np


# import datasets
class Getfiles:
    def __init__(self):
        self.node_file_list = []
        self.node_field = []
        self.edge_file_list = []
        self.edge_field = []
        self.importance_degree = {}
        self.kernel_id = []

    def read(self):
        node_path = './CyberAssetGraphData-main/Dataset/Cleaned_Node.csv'
        edge_path = './CyberAssetGraphData-main/Dataset/Cleaned_Link.csv'

        node_file = pd.read_csv(node_path, header=0, low_memory=False)
        edge_file = pd.read_csv(edge_path, header=0, low_memory=False)
        # print(self.node_file)

        node_file_array = np.array(node_file)
        self.node_file_list = node_file_array.tolist()
        self.node_field = self.node_file_list[0]  # ['id', 'name', 'type', 'industry', 'ID_num']

        edge_file_array = np.array(edge_file)
        self.edge_file_list = edge_file_array.tolist()
        self.edge_field = self.edge_file_list[0]  # ['relation', 'source', 'target']
        # print(node_file_list)
        # print(self.node_field, self.edge_field)
        '''
        for ids in self.node_file_list:
            if ids['is_kernel']:
                self.kernel_id.append()
        '''

    def relation_importance(self):
        # 0 is the most important
        edge_type = [['r_cert', 'r_subdomain', 'r_request_jump', 'r_dns_a'],
                     ['r_whois_name', 'r_whois_email', 'r_whois_phone'],
                     ['r_cert_chain', 'r_cname'],
                     ['r_asn', 'r_cidr']]

        for Type in range(0, len(edge_type)):
            for relation in edge_type[Type]:
                self.importance_degree[relation] = Type
        # print(self.importance_degree)


class Graph:
    def __init__(self):
        self.vertices = []
        self.edge = {}
        self.vex_num = 0
        self.arc_num = 0

    def create(self, vex_list, arc_list):
        self.vex_num = len(vex_list)
        self.arc_num = len(arc_list)
        for vertex in vex_list:
            self.vertices.append(vertex)
        for arc in arc_list:
            self.edge[arc['source']].append([arc['target'], arc['relation']])

    def subgraph(self, node):
        vex_type_num = {}
        edge_type_num = {}
        sub_graph = Graph()
        vis = set()
        queue = []
        level = {}
        for Id in node:
            vis.append(Id)
            queue.append([Id, 3])
            level[Id].append(0)
            sub_graph.vertices.append(Id)
            sub_graph.vex_num += 1
        while queue:
            current_node = queue.pop(0)
            jump = current_node[1]
            current_node = current_node[0]
            vex_type_num[current_node['type']] += 1
            # rule 1 and rule 2
            if level[current_node] < jump:
                for neighbor, relation in self.edge[current_node]:
                    if neighbor not in vis:
                        if file.importance_degree[relation] < 2:
                            queue.append(neighbor, jump-1)
                        elif file.importance_degree[relation] == 2:
                            queue.append(neighbor, min(jump-1, 2))
                        vis.append(neighbor)
                        level[neighbor] = level[current_node] + 1
                        sub_graph.vertices.append(neighbor)
                        sub_graph.edge[current_node].append([neighbor, relation])
                        edge_type_num[relation] += 1
                        sub_graph.vex_num += 1
        return sub_graph

    def asset(self):
        return 0

    def link(self):
        return 0


file = Getfiles()


def preprocess():
    file.read()
    file.relation_importance()


if __name__ == '__main__':
    preprocess()
    g = Graph()

